import uuid
import pymysql  # 替换为PyMySQL
from pymysql import OperationalError, ProgrammingError  # 适配PyMySQL的异常

class HistoryMemory:
    """基于PyMySQL的对话历史管理类（彻底解决参数绑定问题）"""

    def __init__(self, db_config, max_sessions=10, max_turns_per_session=20):
        self.max_sessions = max_sessions
        self.max_turns_per_session = max_turns_per_session
        self.db_config = db_config  # 配置格式不变，无需修改
        self.current_user_id = None
        self.current_session_id = None
        self._create_necessary_tables()  # 初始化表

    def _get_db_connection(self):
        """获取PyMySQL连接（替换mysql.connector）"""
        try:
            # 使用pymysql.connect替换mysql.connector.connect
            conn = pymysql.connect(
                host=self.db_config['host'],
                user=self.db_config['user'],
                password=self.db_config['password'],
                database=self.db_config['database'],
                charset='utf8mb4',
                cursorclass=pymysql.cursors.DictCursor  # 保持字典游标
            )
            return conn
        except OperationalError as e:
            print(f"❌ 数据库连接失败: {str(e)}")
        return None

    def _create_necessary_tables(self):
        """创建表（使用PyMySQL语法）"""
        conn = self._get_db_connection()
        if not conn:
            return
        cursor = conn.cursor()
        try:
            # 会话表（与之前结构一致）
            create_session_sql = """
            CREATE TABLE IF NOT EXISTS chat_sessions (
                id VARCHAR(36) NOT NULL COMMENT '会话ID（UUID）',
                user_id INT NOT NULL COMMENT '关联用户ID',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                PRIMARY KEY (id),
                FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE,
                INDEX idx_user_id (user_id)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
            """
            cursor.execute(create_session_sql)

            # 对话轮次表（与之前结构一致）
            create_turn_sql = """
            CREATE TABLE IF NOT EXISTS chat_turns (
                id VARCHAR(36) NOT NULL COMMENT '轮次ID',
                session_id VARCHAR(36) NOT NULL COMMENT '关联会话ID',
                user_input TEXT NOT NULL,
                agent_response TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                PRIMARY KEY (id),
                FOREIGN KEY (session_id) REFERENCES chat_sessions(id) ON DELETE CASCADE,
                INDEX idx_session_id (session_id)
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
            """
            cursor.execute(create_turn_sql)
            conn.commit()
            print("✅ 表初始化成功（PyMySQL）")
        except ProgrammingError as e:
            conn.rollback()
            print(f"❌ 创建表失败: {str(e)}（表已存在可忽略）")
        finally:
            cursor.close()
            conn.close()

    def _get_or_create_session(self, session_id=None, user_id=None):
        """内部创建会话（逻辑不变，适配PyMySQL）"""
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id)
        except (ValueError, TypeError):
            raise ValueError(f"无效用户ID: {user_id}")
        if not user_id:
            raise ValueError("用户ID不能为空")

        conn = self._get_db_connection()
        if not conn:
            raise Exception("数据库连接失败")

        cursor = conn.cursor()
        try:
            if not session_id:
                session_id = str(uuid.uuid4())

            # 检查会话是否存在
            check_sql = "SELECT id FROM chat_sessions WHERE user_id = %s AND id = %s"
            cursor.execute(check_sql, (user_id, session_id))
            existing_session = cursor.fetchone()

            if not existing_session:
                # 插入新会话
                insert_sql = "INSERT INTO chat_sessions (id, user_id) VALUES (%s, %s)"
                cursor.execute(insert_sql, (session_id, user_id))
                print(f"📥 新建会话: user_id={user_id}, session_id={session_id}")

                # 限制会话数量
                list_sql = "SELECT id FROM chat_sessions WHERE user_id = %s ORDER BY created_at ASC"
                cursor.execute(list_sql, (user_id,))
                all_sessions = cursor.fetchall()
                if len(all_sessions) > self.max_sessions:
                    oldest = all_sessions[0][0]  # PyMySQL返回元组，取第一个元素
                    delete_sql = "DELETE FROM chat_sessions WHERE id = %s"
                    cursor.execute(delete_sql, (oldest,))
                    print(f"🗑️ 删除最早会话: {oldest}")

            self.current_user_id = user_id
            self.current_session_id = session_id
            conn.commit()
            return session_id
        except ProgrammingError as e:
            conn.rollback()
            raise Exception(f"操作会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def start_new_session(self, user_id=None):
        """创建新会话（接口不变）"""
        user_id = user_id or self.current_user_id
        if not user_id:
            raise ValueError("需指定用户ID")
        return self._get_or_create_session(user_id=user_id)

    def add_turn(self, user_input: str, agent_response: str, session_id=None, user_id=None):
        """添加对话轮次（适配PyMySQL）"""
        user_id = user_id or self.current_user_id
        session_id = session_id or self.current_session_id
        if not (user_id and session_id and user_input.strip() and agent_response.strip()):
            raise ValueError("参数不能为空")
        try:
            user_id = int(user_id)
        except (ValueError, TypeError):
            raise ValueError(f"无效用户ID: {user_id}")

        conn = self._get_db_connection()
        if not conn:
            raise Exception("数据库连接失败")

        cursor = conn.cursor()
        try:
            # 检查会话是否存在
            check_sql = "SELECT id FROM chat_sessions WHERE user_id = %s AND id = %s"
            cursor.execute(check_sql, (user_id, session_id))
            if not cursor.fetchone():
                raise Exception(f"会话不存在: {session_id}")

            # 限制轮次数量
            list_sql = "SELECT id FROM chat_turns WHERE session_id = %s ORDER BY created_at ASC"
            cursor.execute(list_sql, (session_id,))
            all_turns = cursor.fetchall()
            if len(all_turns) >= self.max_turns_per_session:
                excess = [turn[0] for turn in all_turns[:-self.max_turns_per_session]]  # 元组取元素
                placeholders = ", ".join(["%s"] * len(excess))
                delete_sql = f"DELETE FROM chat_turns WHERE id IN ({placeholders})"
                cursor.execute(delete_sql, excess)
                print(f"🗑️ 删除{len(excess)}轮旧对话")

            # 插入新对话
            turn_id = str(uuid.uuid4())
            insert_sql = "INSERT INTO chat_turns (id, session_id, user_input, agent_response) VALUES (%s, %s, %s, %s)"
            cursor.execute(insert_sql, (turn_id, session_id, user_input, agent_response))
            conn.commit()
            print(f"✅ 保存对话: turn_id={turn_id[:8]}...")
        except ProgrammingError as e:
            conn.rollback()
            raise Exception(f"保存对话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def get_history(self, session_id=None, user_id=None) -> list:
        """获取对话历史（核心修复：PyMySQL参数绑定）"""
        user_id = user_id or self.current_user_id
        session_id = session_id or self.current_session_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            print(f"⚠️ 无效用户ID: {user_id}")
            return []
        if not (user_id and session_id):
            print(f"⚠️ 参数为空: user_id={user_id}, session_id={session_id}")
            return []

        conn = self._get_db_connection()
        if not conn:
            return []

        cursor = conn.cursor()
        try:
            # 纯单行SQL，参数绑定（PyMySQL无bug）
            # get_history_sql = "SELECT id AS turn_id, user_input AS user, agent_response AS agent, DATE_FORMAT(created_at, '%Y-%m-%d %H:%i:%s') AS timestamp FROM chat_turns WHERE session_id = %s ORDER BY created_at ASC"
            get_history_sql = "SELECT id AS turn_id, user_input AS user, agent_response AS agent, DATE_FORMAT(created_at, '%%Y-%%m-%%d %%H:%%i:%%s') AS timestamp FROM chat_turns WHERE session_id = %s ORDER BY created_at ASC"
            print(f"📤 执行get_history SQL: {get_history_sql} | 参数: ({session_id},)")
            cursor.execute(get_history_sql, (session_id,))  # PyMySQL此处无bug
            history = cursor.fetchall()
            print(f"📥 获取到{len(history)}轮对话")

            return [
                {
                    "id": turn['turn_id'],
                    "timestamp": turn['timestamp'],
                    "user": turn['user'],
                    "agent": turn['agent']
                }
                for turn in history
            ]
        except ProgrammingError as e:
            print(f"❌ 获取历史失败: {str(e)} | SQL: {get_history_sql} | 参数: ({session_id},)")
            return []
        finally:
            cursor.close()
            conn.close()

    def get_sessions_list(self, user_id=None) -> list:
        """获取会话列表（核心修复：PyMySQL参数绑定）"""
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id)
        except (ValueError, TypeError):
            print(f"⚠️ 无效用户ID: {user_id}")
            return []
        if not user_id:
            return []

        conn = self._get_db_connection()
        if not conn:
            return []

        cursor = conn.cursor()
        try:
            # 纯单行SQL，参数绑定（PyMySQL无bug）
            # get_sessions_sql = "SELECT id AS session_id, DATE_FORMAT(created_at, '%Y-%m-%d %H:%i:%s') AS created_at FROM chat_sessions WHERE user_id = %s ORDER BY created_at DESC"
            get_sessions_sql = "SELECT id AS session_id, DATE_FORMAT(created_at, '%%Y-%%m-%%d %%H:%%i:%%s') AS created_at FROM chat_sessions WHERE user_id = %s ORDER BY created_at DESC"
            print(f"📤 执行get_sessions SQL: {get_sessions_sql} | 参数: ({user_id},)")
            cursor.execute(get_sessions_sql, (user_id,))  # PyMySQL此处无bug
            sessions = cursor.fetchall()
            print(f"📥 获取到{len(sessions)}个会话")

            session_list = []
            for session in sessions:
                session_id = session['session_id']
                created_at = session['created_at']
                # 查询会话标题
                get_title_sql = "SELECT user_input FROM chat_turns WHERE session_id = %s ORDER BY created_at ASC LIMIT 1"
                cursor.execute(get_title_sql, (session_id,))
                first_turn = cursor.fetchone()
                title = first_turn['user_input'][:30] + "..." if (first_turn and first_turn['user_input']) else "新会话"
                session_list.append({
                    "id": session_id,
                    "title": title,
                    "timestamp": created_at
                })
            return session_list
        except ProgrammingError as e:
            print(f"❌ 获取会话列表失败: {str(e)} | SQL: {get_sessions_sql} | 参数: ({user_id},)")
            return []
        finally:
            cursor.close()
            conn.close()

    # 以下方法与之前一致，无需修改
    def get_context(self, session_id=None, user_id=None, last_n=5) -> str:
        history = self.get_history(session_id, user_id)
        recent = history[-last_n:] if len(history) >= last_n else history
        parts = []
        for turn in recent:
            parts.append(f"用户: {turn['user']}")
            parts.append(f"助手: {turn['agent']}")
        return "\n".join(parts)

    def clear_session(self, session_id=None, user_id=None):
        user_id = user_id or self.current_user_id
        session_id = session_id or self.current_session_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            return
        if not (user_id and session_id):
            return

        conn = self._get_db_connection()
        if not conn:
            return

        cursor = conn.cursor()
        try:
            cursor.execute("DELETE FROM chat_turns WHERE session_id = %s", (session_id,))
            conn.commit()
            print(f"🧹 清空会话{session_id}：{cursor.rowcount}轮")
        except ProgrammingError as e:
            conn.rollback()
            print(f"❌ 清空会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def delete_session(self, session_id, user_id=None):
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            return
        if not (user_id and session_id):
            return

        conn = self._get_db_connection()
        if not conn:
            return

        cursor = conn.cursor()
        try:
            cursor.execute("DELETE FROM chat_sessions WHERE user_id = %s AND id = %s", (user_id, session_id))
            if cursor.rowcount > 0:
                conn.commit()
                print(f"🗑️ 删除会话{session_id}")
                if self.current_session_id == session_id:
                    self.current_session_id = None
            else:
                print(f"⚠️ 会话不存在: {session_id}")
        except ProgrammingError as e:
            conn.rollback()
            print(f"❌ 删除会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def clear_all_sessions(self, user_id=None):
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            return
        if not user_id:
            return

        conn = self._get_db_connection()
        if not conn:
            return

        cursor = conn.cursor()
        try:
            cursor.execute("DELETE FROM chat_sessions WHERE user_id = %s", (user_id,))
            conn.commit()
            print(f"🗑️ 删除用户{user_id}的{cursor.rowcount}个会话")
            self.current_session_id = None
        except ProgrammingError as e:
            conn.rollback()
            print(f"❌ 清空会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()

    def set_current_session(self, session_id, user_id=None):
        user_id = user_id or self.current_user_id
        try:
            user_id = int(user_id) if user_id else None
        except (ValueError, TypeError):
            return
        if not (user_id and session_id):
            return

        conn = self._get_db_connection()
        if not conn:
            return

        cursor = conn.cursor()
        try:
            cursor.execute("SELECT id FROM chat_sessions WHERE user_id = %s AND id = %s", (user_id, session_id))
            if cursor.fetchone():
                self.current_user_id = user_id
                self.current_session_id = session_id
                print(f"🔄 切换会话: {session_id}")
            else:
                print(f"⚠️ 会话不属于用户: {session_id}")
        except ProgrammingError as e:
            print(f"❌ 切换会话失败: {str(e)}")
        finally:
            cursor.close()
            conn.close()
